package com.zzh.zzhaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.zzh.zzhaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @author: zzh
 * @date: 2025/05/22 20:33:52
 * @version: 1.0
 */
//ToolAgent类 继承 ReActAgent ，具体传递Tool工具，实现具体的功能，实现think方法和act方法
@Data
@EqualsAndHashCode(callSuper = true)
@Slf4j
public class ToolAgent extends ReActAgent {

    //定义可用的工具 (定义final为了注入)
    private final ToolCallback[] availableTools;

    //保存工具调用的返回的结果响应
    private ChatResponse toolCallingResponse;

    //工具调用管理器
    private final ToolCallingManager toolCallingManager;

    //ChatOptions 参数Options
    private final ChatOptions chatOptions;

    public ToolAgent(ToolCallback[] toolCallbacks) {
        this.availableTools = toolCallbacks;
        this.toolCallingManager = ToolCallingManager.builder().build();
        //禁用SpringAI内置的工具调用机制，自己维护选项和消息上下文,因为我们需要控制工具调用，而不是交给AI来全程执行
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)//设置true即可禁用内置的工具调用机制
                .build();
    }

    @Override
    public boolean think() {
        //1.判断是否有下一步执行的提示词，做一个拼接
        if (!StrUtil.isBlank(getNextPrompt())){
            UserMessage userMessage = new UserMessage(getNextPrompt());
            getMessageList().add(userMessage);
        }
        //2.调用AI 获取工具调用的返回结果
        List<Message> messageList = getMessageList();
        Prompt prompt = new Prompt(messageList, chatOptions);//定义Prompt可用让禁止AI托管工具调用！
        try {
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            //记录响应，用于Act
            this.toolCallingResponse = chatResponse;
            //3.解析工具调用结果，获取调用的工具
            //  助手消息
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            //  要调用的工具
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            //  输出日志记录信息
            log.info(getName() + "的思考: " + assistantMessage.getText());
            log.info(getName() + "选择了 " + toolCallList.size() + " 个工具来使用");
            String toolCallInfo = toolCallList.stream()
                    .map(toolCall -> String.format("工具名称：%s，参数：%s",
                            toolCall.name(),
                            toolCall.arguments())
                    )
                    .collect(Collectors.joining("\n"));
            log.info(toolCallInfo);
            if (toolCallList.isEmpty()) {
                // 只有不调用工具时，才记录助手消息
                getMessageList().add(assistantMessage);
                return false;
            } else {
                // 需要调用工具时，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        }catch (Exception e){
            log.error(getName() + "的思考过程遇到了问题: " + e.getMessage());
            getMessageList().add(
                    new AssistantMessage("处理时遇到错误: " + e.getMessage()));
            return false;
        }
    }

    /**
     * 执行工具调用并处理结果
     *
     * @return 执行结果
     */
    @Override
    public String act() {
        if (!toolCallingResponse.hasToolCalls()) {//判断是否有工具调用
            return "没有工具调用";
        }
        // 调用工具
        Prompt prompt = new Prompt(getMessageList(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallingResponse);
        // 记录消息上下文，conversationHistory 已经包含了助手消息和工具调用返回的结果
        setMessageList(toolExecutionResult.conversationHistory());
        // 当前工具调用的结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具 " + response.name() + " 完成了它的任务！结果: " + response.responseData())
                .collect(Collectors.joining("\n"));
        //判断是否调用了终止工具，看调用的工具名称来判断
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream().anyMatch(response -> "doTerminate".equals(response.name()));
        if (terminateToolCalled){
            setState(AgentState.FINISHED);
        }
        log.info(results);
        return results;
    }

}


